package net.xdclass.service.impl;

import com.google.api.gax.core.FixedCredentialsProvider;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.cloud.vision.v1.*;
import com.google.protobuf.ByteString;
import net.xdclass.service.MachineLearningService;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.TakesScreenshot;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.chrome.ChromeDriver;
import org.openqa.selenium.chrome.ChromeOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.scheduling.annotation.Async;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.stereotype.Service;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;

@Service
@EnableAsync
public class MachineLearningServiceImpl implements MachineLearningService {

    private static final String REDIS_CACHE_KEY_PREFIX = "url:malicious:";

    @Autowired
    private ExecutorService executorService;

    @Autowired
    private StringRedisTemplate redisTemplate;

    @Value("${google.cloud.vision.credentials.path}")
    private String googleCloudVisionCredentialsPath;

    static {
        System.setProperty("webdriver.chrome.driver", "path/to/chromedriver");
    }

    @Override
    @Async
    public CompletableFuture<Boolean> detectMac(String url) {
        return CompletableFuture.supplyAsync(() -> {
            String cachedResult = redisTemplate.opsForValue().get(REDIS_CACHE_KEY_PREFIX + url);
            if (cachedResult != null) {
                return Boolean.parseBoolean(cachedResult);
            }

            WebDriver driver = null;
            try {
                ChromeOptions options = new ChromeOptions();
                options.addArguments("--headless");
                options.addArguments("--disable-gpu");
                options.addArguments("--window-size=1920x1080");

                driver = new ChromeDriver(options);
                driver.get(url);

                TakesScreenshot screenshot = (TakesScreenshot) driver;
                byte[] screenshotBytes = screenshot.getScreenshotAs(OutputType.BYTES);

                BufferedImage image = ImageIO.read(new ByteArrayInputStream(screenshotBytes));
                boolean isMalicious = detectMaliciousContent(image);

                redisTemplate.opsForValue().set(REDIS_CACHE_KEY_PREFIX + url, String.valueOf(isMalicious));

                return isMalicious;
            } catch (Exception e) {
                e.printStackTrace();
                return false;
            } finally {
                if (driver != null) {
                    driver.quit();
                }
            }
        }, executorService);
    }

    private boolean detectMaliciousContent(BufferedImage image) {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        try {
            ImageIO.write(image, "png", baos);
        } catch (IOException e) {
            e.printStackTrace();
            return false;
        }
        byte[] imageBytes = baos.toByteArray();

        try (ImageAnnotatorClient visionClient = ImageAnnotatorClient.create(getCredentials())) {
            ByteString imgBytes = ByteString.copyFrom(imageBytes);
            Image img = Image.newBuilder().setContent(imgBytes).build();

            Feature feat = Feature.newBuilder().setType(Feature.Type.SAFE_SEARCH_DETECTION).build();
            AnnotateImageRequest request = AnnotateImageRequest.newBuilder()
                    .addFeatures(feat)
                    .setImage(img)
                    .build();

            BatchAnnotateImagesResponse response = visionClient.batchAnnotateImages(Arrays.asList(request));
            AnnotateImageResponse annotateImageResponse = response.getResponsesList().get(0);

            if (annotateImageResponse.hasError()) {
                System.out.printf("Error: %s%n", annotateImageResponse.getError().getMessage());
                return false;
            }

            SafeSearchAnnotation annotation = annotateImageResponse.getSafeSearchAnnotation();
            if (annotation.getAdult() == Likelihood.LIKELY || annotation.getAdult() == Likelihood.VERY_LIKELY ||
                    annotation.getMedical() == Likelihood.LIKELY || annotation.getMedical() == Likelihood.VERY_LIKELY ||
                    annotation.getRacy() == Likelihood.LIKELY || annotation.getRacy() == Likelihood.VERY_LIKELY ||
                    annotation.getSpoof() == Likelihood.LIKELY || annotation.getSpoof() == Likelihood.VERY_LIKELY ||
                    annotation.getViolence() == Likelihood.LIKELY || annotation.getViolence() == Likelihood.VERY_LIKELY) {
                return true;
            }

            return false;
        } catch (IOException e) {
            e.printStackTrace();
            return false;
        }
    }

    private ImageAnnotatorSettings getCredentials() throws IOException {
        return ImageAnnotatorSettings.newBuilder()
                .setCredentialsProvider(FixedCredentialsProvider.create(GoogleCredentials.fromStream(new FileInputStream(googleCloudVisionCredentialsPath))))
                .build();
    }
}
